# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： crop
Description :
Author : 'li'
date： 2022/6/15
Change Activity:
2022/6/15:
-------------------------------------------------
"""
import numpy as np
import torch
from torch import Tensor

from .... import Dict
from ..transform import ndarray_img_to_pil, pil_to_ndarray
import torchvision.transforms.functional as tf


def crop_image(img, top: int, left: int, height: int, width: int, target=None):
    """
    crop image
    Args:
        img:
        target: elements in target are
        top:
        left:
        height:
        width:

    Returns:

    """
    use_numpy = False
    if isinstance(img, np.ndarray):
        use_numpy = True
        img = ndarray_img_to_pil(img)
    target = target.copy()
    cropped_image = tf.crop(img, top=top, left=left, height=height, width=width)
    if target is None:
        return cropped_image, None
    fields = ["labels", "area", "iscrowd"]
    if "boxes" in target:
        boxes = target["boxes"]
        assert len(boxes.shape) == 2 and boxes.shape[-1] == 4
        max_size = torch.as_tensor([height, width], dtype=torch.float32)
        cropped_boxes = boxes - torch.as_tensor([left, top, left, top])  # boxes shape (boxes size, 4) (xyxy)
        cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size)  # clamp max
        cropped_boxes = cropped_boxes.clamp(min=0)  # clamp min
        area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1)
        target["boxes"] = cropped_boxes.reshape(-1, 4)
        target["area"] = area
        fields.append("boxes")
    if "masks" in target:  # TODO  Check for availability
        # FIXME should we update the area here if there are no boxes?
        target['masks'] = target['masks'][:, top:top + height, left:left + width]
        fields.append("masks")

    # remove elements for which the boxes or masks that have zero area
    if "boxes" in target or "masks" in target:
        # favor boxes selection when defining which elements to keep
        # this is compatible with previous implementation
        if "boxes" in target:
            cropped_boxes = target['boxes'].reshape(-1, 2, 2)
            is_valid_boxes = cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :]
            keep = torch.all(Tensor(is_valid_boxes), dim=1)
        else:
            keep = target['masks'].flatten(1).any(1)

        for field in fields:
            target[field] = target[field][keep]
    if use_numpy:
        return pil_to_ndarray(img), target
    return cropped_image, target
